import argparse
import glob
import json
import os

import pandas as pd

from iflytech_assistant.assistant.dataclasses import FiveSuggestions

user_prompt = """
我现在想学会跟别人高情商沟通。并且有一个要求，需要伪装成{preferenced_style}的人设。

现在我在微信聊天框里面，对面是我的{target}，我想要表达的话是：“{user_input}”

现在我分别找了几个大模型给我建议，他们给了我以下几组表达建议：
{suggestions}

请结合这些结果，从最符合人设，最口语化，最高情商，最能直接在微信键盘端发送，为标准帮我重新组合出5条答案，取各家的优点。 
分析后直接12345输出挑选的结果，不用再给我分析了。

输出格式：
分析：主要分析各个结果的优点，然后结合各个结果的优点，在不改变原始意思的情况下，重新组合出5条答案。
最终结果：
1. xxx
2. xxx
3. xxx
4. xxx
5. xxx
"""

template = {"param": {"max_tokens": 2048}, "id": "h_001", "query": ""}

parser = argparse.ArgumentParser()
parser.add_argument(
    "-n",
    "--name",
    type=str,
)
args = parser.parse_args()

df = pd.read_excel(".vscode/三方大模型试标注.xlsx")

tasks = []
for i, row in df.iterrows():
    user_input = row["input"]
    target = row["target"]
    preferenced_style = row["tags"]

    suggestions = "\n\n".join(
        [str(row["qwen-max"]), str(row["doubao"]), str(row["gpt4o"])]
    )
    up = user_prompt.format(
        user_input=user_input.strip(),
        target=target,
        preferenced_style=preferenced_style,
        suggestions=suggestions,
    )
    query = up

    line = template.copy()
    line["query"] = query
    tasks.append(line)

df = pd.DataFrame(tasks)
# to jsonl
df.to_json(".vscode/refine.jsonl", orient="records", lines=True, force_ascii=False)

labeled_file = ".vscode/e052_输入法_gpt-4o_refine_export_all_lzg_gpt4o_fxx_gpt4o.json"

results = []
if os.path.exists(labeled_file):
    with open(labeled_file, "r") as f:
        for line in f:
            line = json.loads(line)
            answer = line["answer"]
            suggestions = FiveSuggestions.from_llm_response(answer)
            suggestions.suggestions = suggestions.suggestions[-5:]
            results.append(str(suggestions))

    df["refine"] = results
df.to_excel(".vscode/refine.xlsx", index=False)
